import sys
import os
import cv2
import numpy as np
import torch
from PIL import Image, ImageFile
import torch.nn.functional as F
from torch import nn
import copy
from tqdm import tqdm
import pandas as pd

ImageFile.LOAD_TRUNCATED_IMAGES = True
sys.path.append(os.path.abspath('/media/disk/Backup/02drive/13weizhi/unet-pytorch-main'))

from Segmentation.nets.unet import Unet as unet
from utils.utils import cvtColor, preprocess_input, resize_image, show_config


# class unet single
class UnetModel:
    # config dict
    _defaults = {
        "model_path"    : '/media/disk/Backup/02drive/13weizhi/unet-pytorch-main/best_model_Microaneurysms/best_epoch_weights.pth 2',
        "num_classes"   : 2,
        "backbone"      : "resnet50",
        "input_shape"   : [512, 512],
        "mix_type"      : 0,
        "cuda"          : True,
        "onnx"          : False,
    }

    def __init__(self, **kwargs):
        # 初始化配置
        self.__dict__.update(self._defaults)
        self.__dict__.update(kwargs)
        self.cuda = self.cuda and torch.cuda.is_available()

        # 颜色表
        # self._init_colors()

        # 打印配置
        show_config(**self._defaults)

        # 生成网络
        self.net = unet(num_classes=self.num_classes, backbone=self.backbone)
        device = torch.device('cuda:0' if self.cuda else 'cpu')
        self.net.load_state_dict(torch.load(self.model_path, map_location=device))
        self.net.eval()

        if not self.onnx and self.cuda:
            self.net = nn.DataParallel(self.net).cuda()

    # 设置分割区域颜色（暂时没有用上）
    # def _init_colors(self):
    #     if self.num_classes <= 21:
    #         self.colors = [(0, 0, 0), (128, 0, 0), (0, 128, 0)]
    #     else:
    #         hsv = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
    #         rgb = [colorsys.hsv_to_rgb(*x) for x in hsv]
    #         self.colors = [(int(r*255), int(g*255), int(b*255)) for r,g,b in rgb]

    # 检测单个图片
    def detect_image(self, image):
        """检测单张图像，返回掩码mask（非图像）"""
        image = cvtColor(image)
        orininal_h = np.array(image).shape[0]
        orininal_w = np.array(image).shape[1]
        image_data, nw, nh = resize_image(image, (self.input_shape[1], self.input_shape[0]))
        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.cuda()
            pr = self.net(images)[0]
            pr = F.softmax(pr.permute(1, 2, 0), dim=-1).cpu().numpy()
            pr = pr[int((self.input_shape[0] - nh) // 2):int((self.input_shape[0] - nh) // 2 + nh),
                    int((self.input_shape[1] - nw) // 2):int((self.input_shape[1] - nw) // 2 + nw)]
            pr = cv2.resize(pr, (orininal_w, orininal_h), interpolation=cv2.INTER_LINEAR)
            pr = pr.argmax(axis=-1)  # 分类结果

        # 返回mask（True/False数组）
        mask = pr == 1
        return mask


def segment(input_folder, output_folder, model1, model2, model3, csv_file):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    df = pd.read_csv(csv_file)
    # 只处理CSV中记录的图像（采样后的数据）
    img_names = df['file_name'].tolist()

    for img_name in tqdm(img_names, desc="Processing images"):
        img_path = os.path.join(input_folder, img_name)
        try:
            image = Image.open(img_path)

            # 分别用三个模型检测，返回mask
            mask1 = model1.detect_image(image)
            mask2 = model2.detect_image(image)
            mask3 = model3.detect_image(image)

            # 转换原图为numpy数组
            result_np = np.array(image)

            # 合并三个mask（任一模型检测到的区域都保留）
            combined_mask = mask1 | mask2 | mask3

            # 只保留mask区域的像素，其他区域设为黑色
            result_np[~combined_mask] = [0, 0, 0]

            # 保存结果
            result_image = Image.fromarray(result_np.astype('uint8'))
            save_path = os.path.join(output_folder, img_name)
            result_image.save(save_path)
            print(f"{img_name}保存到{save_path}")

        except (OSError, AttributeError) as e:
            print(f"图片处理失误{img_name}: {e}")
            idx_to_remove = img_name.split('.')[0]
            df = df[df['idx'] != idx_to_remove]
            print(f"删除{idx_to_remove}")

    df.to_csv(csv_file, index=False)
    print(f"CSV已更新: {csv_file}")


if __name__ == "__main__":
    input_folder = "path/to/input_images"
    output_folder = "path/to/output_images"
    csv_file = "path/to/data.csv"

    # 加载三个模型，分别对应三个病理特征
    model1 = UnetModel(
        model_path="path/to/model1_weights.pth",
    )

    model2 = UnetModel(
        model_path="path/to/model2_weights.pth",
    )

    model3 = UnetModel(
        model_path="path/to/model3_weights.pth",
    )

    # 执行推理
    segment(input_folder, output_folder, model1, model2, model3, csv_file)